package io.openkunlun.javadsl;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class Context {

    public static final Context ROOT = new Context(null);

    private final Context parent;
    private ConcurrentMap<Object, Object> data;
    private ConcurrentMap<Object, Object> localData;

    /**
     * @param parent
     */
    Context(Context parent) {
        this.parent = parent;
    }

    /**
     * Get some data from the context.
     *
     * @param key the key of the data
     * @param <T> the type of the data
     * @return the data
     */
    @SuppressWarnings("unchecked")
    public <T> T get(Object key) {
        return (T) contextData().get(key);
    }

    /**
     * Put some data in the context.
     * <p>
     * This can be used to share data between different handlers that share a context
     *
     * @param key   the key of the data
     * @param value the data
     */
    public void put(Object key, Object value) {
        contextData().put(key, value);
    }

    /**
     * Remove some data from the context.
     *
     * @param key the key to remove
     * @return true if removed successfully, false otherwise
     */
    public boolean remove(Object key) {
        return contextData().remove(key) != null;
    }

    /**
     * Get some local data from the context.
     *
     * @param key the key of the data
     * @param <T> the type of the data
     * @return the data
     */
    @SuppressWarnings("unchecked")
    public <T> T getLocal(Object key) {
        return (T) localContextData().get(key);
    }

    /**
     * Put some local data in the context.
     * <p>
     * This can be used to share data between different handlers that share a context
     *
     * @param key   the key of the data
     * @param value the data
     */
    public void putLocal(Object key, Object value) {
        localContextData().put(key, value);
    }

    /**
     * Remove some local data from the context.
     *
     * @param key the key to remove
     * @return true if removed successfully, false otherwise
     */
    public boolean removeLocal(Object key) {
        return localContextData().remove(key) != null;
    }

    public Context attach() {
        Context prev = storage().attach(this);
        return prev == null ? ROOT : prev;
    }

    public void detach(Context toAttach) {
        storage().detach(this, toAttach);
    }

    /**
     * @return
     */
    public boolean isCurrent() {
        return current() == this;
    }

    /**
     * @return
     */
    public Context duplicate() {
        return new Context(this);
    }

    /**
     * @param r
     */
    public void run(Runnable r) {
        Context previous = attach();
        try {
            r.run();
        } finally {
            detach(previous);
        }
    }

    /**
     * @param c
     * @param <V>
     * @return
     * @throws Exception
     */
    public <V> V call(Callable<V> c) throws Exception {
        Context previous = attach();
        try {
            return c.call();
        } finally {
            detach(previous);
        }
    }

    private synchronized ConcurrentMap<Object, Object> contextData() {
        if (null != parent) {
            return parent.contextData();
        }

        if (null == data) {
            data = new ConcurrentHashMap<>();
        }
        return data;
    }

    private synchronized ConcurrentMap<Object, Object> localContextData() {
        if (null == localData) {
            localData = new ConcurrentHashMap<>();
        }

        return localData;
    }

    /**
     * @return
     */
    public static Context fork() {
        return current().duplicate();
    }

    /**
     * @return
     */
    public static Context current() {
        Context current = storage().current();
        return current == null ? ROOT : current;
    }

    private static ContextStorage storage() {
        return LazyStorage.storage;
    }

    private static final class LazyStorage {
        static final ContextStorage storage;

        static {
            // todo Supported More Storage
            storage = new ThreadLocalStorage();
        }
    }

    private static final class ThreadLocalStorage implements ContextStorage {

        private static final ThreadLocal<Context> LOCAL_CONTEXT = new ThreadLocal<>();
        private static final Logger LOG = LoggerFactory.getLogger(ThreadLocalStorage.class);

        @Override
        public Context attach(Context toAttach) {
            Context current = current();
            LOCAL_CONTEXT.set(toAttach);
            return current;
        }

        @Override
        public void detach(Context toDetach, Context toRestore) {
            if (current() != toDetach) {
                LOG.error("Context was not attached when detaching", new Throwable().fillInStackTrace());
            }

            if (toRestore != Context.ROOT) {
                LOCAL_CONTEXT.set(toRestore);
            } else {
                LOCAL_CONTEXT.set(null);
            }
        }

        @Override
        public Context current() {
            Context current = LOCAL_CONTEXT.get();
            return current == null ? Context.ROOT : current;
        }
    }
}
